#-*- coding:utf-8 -*-
import cv2,os,time,copy
from config import Config
import numpy as np

config = Config()
COLOR_WHITE = (255,255,255)
COLOR_BLACK = (0,0,0)
COLOR_BORDER = (0,0,255)
threathhold=0.99
strokeColorBg= COLOR_BLACK
strokeColorFg= COLOR_WHITE


def MSE(image,gt):
      if image.max()>1:
         image=image/255.0
      if gt.max()>1:
         gt=gt/255.0
      if len(image.shape)>2:
           image=image[:,:,0]
      if len(gt.shape)>2:
           gt=gt[:,:,0]
  
      width=image.shape[0]
      height=image.shape[1]
      sumup=0
      for j in range(width):
         for k in range(height):
             sumup=sumup+(np.float32(image[j,k])-np.float32(gt[j,k]))**2
      MSE=(sumup+0.0)/(width*height)
      return MSE


def genStrokesForOne(alphaImg, rgbImg):
      img_width = alphaImg.shape[0]
      img_height = alphaImg.shape[1]
      num_pt_bg =0
      num_pt_fg =0
      for i in range(img_width):
          for j in range(img_height):
              if alphaImg[i,j]==0:
                    num_pt_bg+=1
              elif alphaImg[i,j]==255:
                     num_pt_fg+=1
      num_pt_all=img_height*img_width
      Bgratio = num_pt_bg/num_pt_all
      Fgratio = num_pt_fg/num_pt_all
      if Bgratio >threathhold:
            rgbImg[:,:]=strokeColorBg
      elif Fgratio >threathhold:
            rgbImg[:,:]=strokeColorFg
      return rgbImg.copy()


def genStroke(strokeMap, loc, rgbImg, alphaImg,window_size=[400*0.05,400*0.05],im_sz=400):
     leftup = [max(0,(loc[0] - window_size[0]/2)),max(0,(loc[1] - window_size[1]/2))]    
     rightdown = [min(im_sz,(loc[0] + window_size[0]/2)),min(im_sz,(loc[1] + window_size[1]/2))]
          
     temp_alpha_img = alphaImg[leftup[1]:rightdown[1], leftup[0]:rightdown[0]]
     temp_rgb_img = rgbImg[leftup[1]:rightdown[1],leftup[0]:rightdown[0]]
          
     strokes_img = genStrokesForOne(temp_alpha_img,temp_rgb_img)
     strokeMap[leftup[1]:rightdown[1],leftup[0]:rightdown[0]] = strokes_img
     return strokeMap






def manualDraw(img,winLoc,winSize=400*0.05):
  #print winLoc
  def draw_call(event,x,y,flags,param):
    global img2,cmode
    winSize_half = winSize/2.0
    leftup = tuple((winLoc-winSize_half).astype(np.int16).tolist())
    rightdown = tuple((winLoc+winSize_half).astype(np.int16).tolist())
    if event == cv2.EVENT_LBUTTONDOWN:
      if cmode == True:  #前景
        cv2.rectangle(img2,leftup,rightdown,COLOR_WHITE,-1)
      else:  #背景
        cv2.rectangle(img2,leftup,rightdown,COLOR_BLACK,-1)
    
  def img_with_shadow(img2,winLoc,winSize):
    img = img2.copy()
    winSize_half = winSize/2.0
    leftup = tuple((winLoc-winSize_half).astype(np.int16).tolist())
    rightdown = tuple((winLoc+winSize_half).astype(np.int16).tolist())
    cv2.rectangle(img,leftup,rightdown,COLOR_BORDER,2)
    return img

  global img2,drawing,cmode,ix,iy,jx,jy
  img2 = img.copy()
  drawing = False
  cmode = False
  ix = iy = jx = jy = -1
  cv2.namedWindow('Image')
  cv2.setMouseCallback('Image',draw_call)
  while(True):
    img3 = img_with_shadow(img2,winLoc,winSize)
    cv2.imshow('Image',img3)
    k = cv2.waitKey(1)&0xFF
    if k == ord('f'): #前景
      cmode = True
    elif k == ord('b'): #背景
      cmode = False
    elif k == ord('n'): #保存
      img = img2
    elif k == 27:     #ESC
      break
    time.sleep(0.05)
  cv2.destroyAllWindows()
  return img

def findRegionbyOrder(loc_array,stroke):
       loc=loc_array[0]
       strokeMap=copy.deepcopy(stroke)
       strokeMap=genStroke(strokeMap,loc,image,gt)
       cv2.imwrite('stroke.png',strokeMap)
       systemCall='./SharedMatting -i ./image.png -t ./stroke.png -a'
       os.system(systemCall)
       time.sleep(0.05)
       matte=cv2.imread('shared_matting_alpha.bmp')
       os.system('rm shared_matting_alpha.bmp')
       os.system('rm stroke.png')
       cv2.imwrite('./matte.png',matte)
       matte=cv2.imread('./matte.png',0)
       mse=MSE(matte,gt)
       return loc,mse,strokeMap

def findRegionRandom(loc_array,stroke):
       loc=loc_array[0]
       strokeMap=copy.deepcopy(stroke)
       strokeMap=genStroke(strokeMap,loc,image,gt)
       cv2.imwrite('stroke.png',strokeMap)
       systemCall='./SharedMatting -i ./image.png -t ./stroke.png -a'
       os.system(systemCall)
       time.sleep(0.05)
       matte=cv2.imread('shared_matting_alpha.bmp')
       os.system('rm shared_matting_alpha.bmp')
       os.system('rm stroke.png')
       cv2.imwrite('./matte.png',matte)
       matte=cv2.imread('./matte.png',0)
       mse=MSE(matte,gt)
       return loc,mse,strokeMap


def findRegion(loc_arry,stroke,Min_mse=-1):
    for loc in loc_array:
        strokeMap=copy.deepcopy(stroke)
        strokeMap=genStroke(strokeMap,loc,image,gt)
        cv2.imwrite('stroke.png',strokeMap)
        systemCall='./SharedMatting -i ./image.png -t ./stroke.png -a'
        os.system(systemCall)
        time.sleep(0.05)
        matte=cv2.imread('shared_matting_alpha.bmp')
        os.system('rm shared_matting_alpha.bmp')
        os.system('rm stroke.png')
        cv2.imwrite('./matte.png',matte)
        matte=cv2.imread('./matte.png',0)
        mse=MSE(matte,gt)
        if Min_mse<0:
            Min_mse=mse
            best_loc=loc
            bestStroke=strokeMap
        if Min_mse>=mse:
           best_loc=loc
           Min_mse=mse
           bestStroke=strokeMap
    return loc,Min_mse,bestStroke





if __name__=='__main__':
   locX_file=open('orderX.txt')
   locY_file=open('orderY.txt')
   image=cv2.imread('image.png')
   gt=cv2.imread('gt.png',0)
 
   
   locX=locX_file.readline()
   locY=locY_file.readline()
   locX=float(locX)
   locY=float(locY)
   loc=np.array((locX,locY))
   loc_array=loc
   while 1:
       locX=locX_file.readline()
       locY=locY_file.readline()
       if not locX or not locY:
          break
       locX=float(locX)
       locY=float(locY)
       loc=np.array((locX,locY))
       loc_array=np.vstack((loc_array,loc))
       #stroke=np.zeros((config.img_sz, config.img_sz, config.strokeMap_channels))
       #stroke=manualDraw(image,np.array([locX,locY]))
       #cv2.imshow('stroke',stroke)
       #cv2.waitKey()
       #cv2.destroyAllWindows()
       
   locX_file.close()
   locY_file.close()
   #print loc_array.shape
   #np.random.shuffle(loc_array)
   #print '###################'
   #print loc_array

   stroke=image
   loc_list=[]
   mse_list=[]

   #random shuffle loc_array
   np.random.shuffle(loc_array)

   while(loc_array.size):
     #loc,mse,stroke=findRegion(loc_array,stroke)
     loc,mse,stroke=findRegionbyOrder(loc_array,stroke)
     loc_list.append(loc)
     mse_list.append(mse)
     loc_array=filter(lambda thisloc: not np.equal(loc,thisloc).all(),loc_array)
     loc_array=np.array(loc_array)
     print '######################'
     print loc_array
     print 'size is', loc_array.size
     print '######################'
   print loc_list
   print '######################3'
   print  mse_list
       
       
            
            
            
            

            
    








